/*
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
/*
 * @author Anthony Surma
 */


import java.io.*;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Scanner;

public class XShape
{    
    final static String   VERSION    = "0.13b";
    final static String   CONFIG     = "XSHAPE.config";
    final static String   EXE        = "Binaries\\Win32\\XComGame.exe"; 
    final static String   TOC        = "XComGame\\PCConsoleTOC.txt";
    final static String   CFOLDER    = "XComGame\\CookedPCConsole\\";    
    final static String   COMPRESSED = ".uncompressed_size";
    final static String   PERIOD     = "\\.";    
    final static int      SEPERATOR  = 0x00;
    final static int      BUFFERPAD  = 100;
          static boolean  isVerbose  = false;
          static long     startAtByteNum;
          static long     numMaxBytes;
          static String[] filesList;
          static byte[][] hashes;
    
    class Data
    {
        long address;
        int fileNum;
        
        Data(long address, int fileNum)
        {
            this.address = address;
            this.fileNum = fileNum;
        }
    }
          
    public static void main(String args[]) throws Exception
    {
        XShape xshape = new XShape();                
        System.out.println("XCOM SHA Patcher for Executable (XSHAPE) v" + VERSION);
        int exitCode;
        
        if ((exitCode = runSanityChecks(args)) == 0) 
        {
            xshape.calculateHashes();
            xshape.startPatching();
        }   
        else showUsage(exitCode);
                
        System.exit(exitCode);         
    }
    
    void calculateHashes() throws NoSuchAlgorithmException, FileNotFoundException, IOException
    {     
        MessageDigest md = MessageDigest.getInstance("SHA"); 
        byte[] dataBytes = new byte[8096];                    
        hashes = new byte[filesList.length][20];
        
        for (int i = filesList.length - 1; i >= 0; --i)
        {
            FileInputStream fis = new FileInputStream(new File(CFOLDER + filesList[i]));           
            int numRead = 0;             
            
            while ((numRead = fis.read(dataBytes)) > 0) md.update(dataBytes, 0, numRead);                                              
            
            hashes[i] = md.digest();  
            md.reset();          
            fis.close();
        }                        
    }
        
    void startPatching() throws FileNotFoundException, IOException
    {                          
        ArrayList<Data> results = new ArrayList<>(filesList.length);
        
        RandomAccessFile raf = new RandomAccessFile(EXE, "rw");                
        final long max = numMaxBytes;
        byte[] buffer  = new byte[(int)(numMaxBytes - startAtByteNum + BUFFERPAD)];      
        byte[][] filenames = new byte[filesList.length][];
        int numFound = 0;
        long current;
        
        for (int i = buffer.length - 1, j = BUFFERPAD; j > 0; --j, --i) buffer[i] = 0x11;
                        
        vPrintln("\nFound " + numMaxBytes + " bytes in " + EXE);
        vPrintln("Starting at byte number: " + startAtByteNum);        
        vPrintln("Buffer size: " + (buffer.length - BUFFERPAD) + "\n");    
        
        raf.seek(startAtByteNum);
        current = raf.getFilePointer(); 
        
        for (int i = filesList.length - 1; i >= 0; --i) filenames[i] = filesList[i].getBytes();
        for (int i = buffer.length - 1, j = BUFFERPAD; j > 0; --j, --i) buffer[i] = 0x11;                                     
                
        raf.read(buffer);

        for (int curFile = 0; curFile < filesList.length; ++curFile)
        {
            byte[] filename = filenames[curFile];                        
            final int fileLength = filename.length;
            final int end = buffer.length - BUFFERPAD;
            int cur = 0;
           
            while (cur < end) 
            {
                int i = cur;
                int j = 0;

                while (j < fileLength)
                {
                    if (filename[j] == buffer[i])
                    {
                        ++j;   
                        ++i;
                    }
                    else
                    {
                        break;
                    }
                }

                if ((j == fileLength) && (buffer[i] == SEPERATOR))
                {                        

                    byte[] foundHash = Arrays.copyOfRange(buffer, i + 1, i + 21);
                    vPrintln();
                    vPrintln(new String(filename) + " entry found at: " + (current + i - j) + 
                                                    "\nFound Hash.: " + vHashToString(foundHash));                          
                    vPrintln("Actual Hash: " + vHashToString(hashes[curFile]));    
                    if (!Arrays.equals(hashes[curFile], foundHash)) results.add(new Data(current + i + 1, curFile));                                                                      
                    ++numFound;                                         
                    break;                                           
                }

                if (numFound == filesList.length) break; 
                
                ++cur;
            }                                                                                                              
        }
                                        
        for (Data d : results)
        {            
            raf.seek(d.address);
            raf.write(hashes[d.fileNum]);
            vPrintln(filesList[d.fileNum] + " SHA hash updated.");
        }
        
        raf.close();                   
        
        if (numFound != filesList.length) { showUsage(16); System.exit(16);}    
        else if (results.isEmpty()) System.out.println("\nNo changes needed to be made.\n");        
        else System.out.println(results.size() + " SHA hash(es) updated.\n");                
    }
    
    static int runSanityChecks(String[] args) throws IOException
    {        
        File configFile = new File(CONFIG);
        File exeFile = new File(EXE);
        Scanner scan;
        int count = 0;
        
        if (args.length == 0) return 1;        
        if (!configFile.canRead()) return 2;        
        if (configFile.length() == 0) return 3;
                                 
        scan = new Scanner(configFile);
        while (scan.hasNextLine())
        {                        
            String line = scan.nextLine().trim();            
            if (line.isEmpty()) 
            {
                scan.close();
                return 4;
            }
            if (!isCompressed(line)) ++count;            
        }
        scan.close();  
        
        if (count == 0) return 5;        
        filesList = new String[count];        
        scan = new Scanner(configFile);
        for (int i = 0; i < count; i++) // ugly haha
        {            
            String line = scan.nextLine().trim().toLowerCase();
            if (isCompressed(line))
            {
                --i;
            }
            else
            {
                filesList[i] = line;                                                            
            }                        
        }
        scan.close();
                
        if (args[0].contains("-"))
        {
            if (args.length != 2) return 6;                        
            if (args[0].contains("v")) isVerbose = true;    
            
            try 
            { 
                startAtByteNum = new Long(args[1]).longValue(); 
            } catch(NumberFormatException e) { return 7; };
        }        
        else 
        {
            if (args.length != 1) return 8;
            startAtByteNum = new Long(args[0]).longValue();
        }
        
        vPrintln();
        for (String s : filesList)
        {                 
            boolean valid = false;
            vPrintln(s);
            String[] split = s.split(PERIOD);            
            if (split.length != 2) return 9;                                  
            if (split[1].length() != 3) return 10;            
            
            scan = new Scanner(new File(TOC));
            while (scan.hasNextLine())
            {
                String line = scan.nextLine().trim().toLowerCase();
                if (line.contains(s))
                {                   
                    valid = true;
                    break;
                }                                   
            }
            if (!valid) return 11;            
        }

        if (startAtByteNum < 0) return 12;        
        if (!exeFile.canRead()) return 13;        
        if (!exeFile.canWrite()) return 14;
        
        numMaxBytes = exeFile.length();        
        if (startAtByteNum >= numMaxBytes) return 15;
                                          
        return 0; // everything is fine
    }
    
    static void showUsage(int errorCode)
    {
        System.out.println("\nERROR(" + errorCode + ")\nUSAGE: java -jar XSHAPE.jar -<verbose> <starting byte>\n" + 
                           "More information: http://code.google.com/p/xshape \n"); 
    }
            
    static boolean isCompressed(String s) { return new File(CFOLDER + s + COMPRESSED).exists(); }
    
    static String vHashToString(byte[] b)
    {
        StringBuilder result = new StringBuilder(20);
        if (isVerbose)
        {
            for (int i = 0; i < b.length; ++i) 
            {                             
                String hexString = Integer.toHexString((int)b[i] & 0xff).toUpperCase();           
                if (hexString.length() == 1) hexString = "0" + hexString;  
                result.append(hexString).append(" ");
            }           
        }
        return result.toString();
    }
    
    static void vPrint(Object o) { if (isVerbose) System.out.print(o.toString()); }
    static void vPrintln(Object o) { if (isVerbose) System.out.println(o.toString()); }
    static void vPrintln() { if (isVerbose) System.out.println(); }    
}